// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/utils.h>
#include <pollux/common/utils.h>
#include <pollux/expression/vector_writers.h>
#include <pollux/vector/nulls_builder.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {
    namespace generator_spec_utils {
        vector_size_t getRandomIndex(FuzzerGenerator &rng, vector_size_t maxIndex) {
            std::uniform_int_distribution<vector_size_t> indexGenerator(
                0, maxIndex); // generates index in [0, maxIndex]
            return indexGenerator(rng);
        }

        BufferPtr generateNullsBuffer(
            FuzzerGenerator &rng,
            memory::MemoryPool *pool,
            vector_size_t vectorSize,
            double nullProbability) {
            NullsBuilder builder{vectorSize, pool};
            for (size_t i = 0; i < vectorSize; ++i) {
                if (fuzzer::coinToss(rng, nullProbability)) {
                    builder.set_null(i);
                }
            }
            return builder.build();
        }

        BufferPtr generateIndicesBuffer(
            FuzzerGenerator &rng,
            memory::MemoryPool *pool,
            vector_size_t bufferSize,
            vector_size_t baseVectorSize) {
            BufferPtr indices = AlignedBuffer::allocate<vector_size_t>(bufferSize, pool);
            auto rawIndices = indices->asMutable<vector_size_t>();
            auto indicesGenerator =
                    std::uniform_int_distribution<vector_size_t>(0, baseVectorSize - 1);

            for (size_t i = 0; i < bufferSize; ++i) {
                rawIndices[i] = indicesGenerator(rng);
            }
            return indices;
        }
    } // namespace generator_spec_utils

    namespace fuzzer {
        template<>
        void writeOne<TypeKind::VARCHAR>(
            const variant &v,
            exec::GenericWriter &writer) {
            writer.template castTo<Varchar>() = v.value<TypeKind::VARCHAR>();
        }

        template<>
        void writeOne<TypeKind::VARBINARY>(
            const variant &v,
            exec::GenericWriter &writer) {
            writer.template castTo<Varbinary>() = v.value<TypeKind::VARBINARY>();
        }

        template<>
        void writeOne<TypeKind::ARRAY>(const variant &v, exec::GenericWriter &writer) {
            auto &writerTyped = writer.template castTo<Array<Any> >();
            const auto &elements = v.array();
            for (const auto &element: elements) {
                if (element.isNull()) {
                    writerTyped.add_null();
                } else {
                    POLLUX_DYNAMIC_TYPE_DISPATCH(
                        writeOne, element.kind(), element, writerTyped.add_item());
                }
            }
        }

        template<>
        void writeOne<TypeKind::MAP>(const variant &v, exec::GenericWriter &writer) {
            auto &writerTyped = writer.template castTo<Map<Any, Any> >();
            const auto &map = v.map();
            for (const auto &pair: map) {
                const auto &key = pair.first;
                const auto &value = pair.second;
                POLLUX_CHECK(!key.isNull());
                if (value.isNull()) {
                    POLLUX_DYNAMIC_TYPE_DISPATCH(
                        writeOne, key.kind(), key, writerTyped.add_null());
                } else {
                    auto writers = writerTyped.add_item();
                    POLLUX_DYNAMIC_TYPE_DISPATCH(
                        writeOne, key.kind(), key, std::get<0>(writers));
                    POLLUX_DYNAMIC_TYPE_DISPATCH(
                        writeOne, value.kind(), value, std::get<1>(writers));
                }
            }
        }

        template<>
        void writeOne<TypeKind::ROW>(const variant &v, exec::GenericWriter &writer) {
            auto &writerTyped = writer.template castTo<DynamicRow>();
            const auto &elements = v.row();
            column_index_t i = 0;
            for (const auto &element: elements) {
                if (element.isNull()) {
                    writerTyped.set_null_at(i);
                } else {
                    POLLUX_DYNAMIC_TYPE_DISPATCH(
                        writeOne, element.kind(), element, writerTyped.get_writer_at(i));
                }
                i++;
            }
        }
    } // namespace fuzzer
} // namespace kumo::pollux
